package org.example.lanchain.config;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import org.example.lanchain.store.MongoChatMemoryStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.Arrays;
import java.util.List;

@Configuration
public class QwMemoryChat {

    @Autowired
    MongoChatMemoryStore mongoChatMemoryStore;

    @Autowired
    private EmbeddingModel embeddingModel;

    @Autowired
    private EmbeddingStore embeddingStore;

    @Bean
    public ChatMemory qwChatMemory() {
        return MessageWindowChatMemory.withMaxMessages(10);
    }

    @Bean
    public ChatMemoryProvider chatMemoryProvider() {
        return memoryId -> MessageWindowChatMemory.builder()
                .id(memoryId)
                .maxMessages(20) // 记忆长度
                .chatMemoryStore(mongoChatMemoryStore)
                .build();
    }

    // @Bean
    // ContentRetriever contentRetrieverXiaozhi() {
    //     // 加载文档 解析文档
    //     Document document1 = FileSystemDocumentLoader.loadDocument("D:\\BaiduNetdiskDownload\\knowledge\\knowledge\\医院信息.md");
    //     Document document2 = FileSystemDocumentLoader.loadDocument("D:\\BaiduNetdiskDownload\\knowledge\\knowledge\\科室信息.md");
    //     Document document3 = FileSystemDocumentLoader.loadDocument("D:\\BaiduNetdiskDownload\\knowledge\\knowledge\\神经内科.md");
    //     List<Document> documents = Arrays.asList(document1, document2, document3);
    //
    //     // 内存向量存储
    //     InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
    //     // 使用默认的文档分割器，并将分割后的文本向量化存到embeddingStore中
    //     EmbeddingStoreIngestor.ingest(documents, embeddingStore);
    //     // 从embeddingStore检索和查询相关内容的信息
    //     return EmbeddingStoreContentRetriever.from(embeddingStore);
    // }

    @Bean
    ContentRetriever contentRetrieverXiaozhiPinecone() {
        return EmbeddingStoreContentRetriever
                .builder()
                .embeddingModel(embeddingModel)
                .embeddingStore(embeddingStore)
                // 最大匹配数
                .maxResults(1)
                // 最小相似度
                .minScore(0.8)
                .build();
    }

}
